package com.huhx.demo.lock;

import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

import org.I0Itec.zkclient.IZkDataListener;
import org.I0Itec.zkclient.ZkClient;
import org.I0Itec.zkclient.exception.ZkException;
import org.I0Itec.zkclient.exception.ZkNoNodeException;

public class ZkDistributedLock {
	private static final String LOCKPATH = "/lock";
	private static final String SEQLOCKPATH = "/seqLock";
	
	private String outPath;

	private ZkClient client;
	
	public ZkDistributedLock(String hosts) {
		client = new ZkClient(hosts);
	}
	
	/**
	 * 竞争锁
	 * @param lockName 锁名称
	 * @param waitTime 获取锁超时时间（毫秒）
	 * @return
	 * @throws InterruptedException
	 */
	public boolean lock(String lockName, long waitTime) throws InterruptedException {
		final long beginMillis = System.currentTimeMillis();
		long millis;
		String path = LOCKPATH.concat("/").concat(lockName);
		while (true) {
			if (client.exists(path)) {
				final CountDownLatch latch = new CountDownLatch(1);
				IZkDataListener listener = new IZkDataListener() {
					//节点数据变化回调
					public void handleDataChange(String arg0, Object arg1) throws Exception {
					}
					//节点被删除回调
					public void handleDataDeleted(String arg0) throws Exception {
						latch.countDown();
					}
				};
				try {
					//如果节点不存在会出现异常
					client.subscribeDataChanges(path, listener);
					millis = System.currentTimeMillis() - beginMillis;
					if (millis >= waitTime) {
						return false;
					}
					boolean b = latch.await(waitTime - millis, TimeUnit.MILLISECONDS);
					if (!b) {
						return false;
					}
				} catch (ZkNoNodeException e) {
					
				} finally {
					client.unsubscribeDataChanges(path, listener);
				}
			}
			//尝试创建临时节点
			try {
				client.createEphemeral(path); //节点已经存在时，会报异常
				return true;
			} catch (ZkException e) {
				millis = System.currentTimeMillis() - beginMillis;
				if (millis >= waitTime) {
					return false;
				}
			}
		}
	}
	
	/**
	 * 顺序锁
	 * @param lockName 锁名称
	 * @param waitTime 获取锁超时时间（毫秒）
	 * @return
	 * @throws InterruptedException
	 */
	public boolean seqLock(String lockName, long waitTime) throws InterruptedException {
		final long beginMillis = System.currentTimeMillis();
		long millis;
		String path = SEQLOCKPATH.concat("/").concat(lockName);
		//创建临时顺序节点
		outPath = client.createEphemeralSequential(path, null);
		//System.out.println("lock: " + Thread.currentThread().getName() + " lockPaht: " + outPath);
		while (true) {
			//获取所有临时顺序节点
			List<String> children = client.getChildren(SEQLOCKPATH);
			Collections.sort(children); //从小到大排序
			String sequenceNodeName = outPath.substring(SEQLOCKPATH.length() + 1);
			//获取刚才创建的顺序节点在lock的所有子节点中排序位置，如果是排序为0，则表示获取到了锁
			int idx = children.indexOf(sequenceNodeName);
			if (idx == 0) {
				return true;
			}
			//如果在子节点中没有找到刚创建的临时节点，可能是因为网络闪断导致的，此处一般需重试下
			if (idx < 0) {
				millis = System.currentTimeMillis() - beginMillis;
				if (millis >= waitTime) {
					return false;
				} else {
					continue;
				}
			}
			//监听子节点中比自己小的节点
			String watchPath = SEQLOCKPATH.concat("/").concat(children.get(idx - 1));
			final CountDownLatch latch = new CountDownLatch(1);
			IZkDataListener listener = new IZkDataListener() {
				//节点数据变化回调
				public void handleDataChange(String paramString, Object paramObject) throws Exception {
				}
				//节点被删除回调
				public void handleDataDeleted(String paramString) throws Exception {
					latch.countDown();
				}
			};
			
			try {
				//如果节点不存在会出现异常
				client.subscribeDataChanges(watchPath, listener);
				millis = System.currentTimeMillis() - beginMillis;
				if (millis >= waitTime) {
					return false;
				}
				boolean b = latch.await(waitTime - millis, TimeUnit.MILLISECONDS);
				if (!b) {
					return b;
				}
			} catch (ZkNoNodeException e) {
				millis = System.currentTimeMillis() - beginMillis;
				if (millis >= waitTime) {
					return false;
				}
			} finally {
				client.unsubscribeDataChanges(watchPath, listener);
			}
		}
	}
	
	public void unlock(String lockName) {
		String path = LOCKPATH.concat("/").concat(lockName);
		client.delete(path);
	}
	
	public void unSeqLock() {
		client.delete(outPath);
	}
	
	public void close() {
		client.close();
	}
}
